package com.ds.infra.mq.core.storage;

import com.ds.infra.mq.core.config.StorageConfig;
import com.ds.infra.mq.core.protocol.Message;
import com.google.common.collect.Maps;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * 作用：内存消息存储实现，用于临时测试
 *
 * @author WeiShaoying
 */
@Slf4j
public class MemoryMessageStore {

    private final StorageConfig config;

    // 存储结构：topic-partition -> 消息列表(最多保留1万条消息，过期自动清理)
    private final ConcurrentMap<String, List<Message>> topicMessages = Maps.newConcurrentMap();

    /**
     * 消费偏移量：key=consumerId:topic , value=offset
     * 保留不同消费者的消费进度
     */
    private final ConcurrentMap<String, Long> consumerOffsets = Maps.newConcurrentMap();

    public MemoryMessageStore(StorageConfig config) {
        this.config = config;
        log.info("Initialized MemoryMessageStore with config: {}", config);
    }

    /**
     * 存储消息到指定主题（默认分区）
     */
    public void putMessage(String topic, Message message) {
        putMessage(topic, 0, message); // 默认分区0
    }

    /**
     * 存储消息到指定Topic分区
     */
    public void putMessage(String topic, int partition, Message message) {
        String partitionKey = getPartitionKey(topic, partition);
        topicMessages.computeIfAbsent(partitionKey, k -> new CopyOnWriteArrayList<>())
                .add(message);

        log.debug("Stored message to {} (size={})", partitionKey, topicMessages.get(partitionKey).size());

        cleanExpiredMessages(partitionKey);
    }

    /**
     * 从指定主题分区获取消息（基于消费者偏移量）
     */
    public Message getMessage(String topic, int partition, String consumerId) {
        String partitionKey = getPartitionKey(topic, partition);
        String consumerKey = getConsumerKey(consumerId, topic, partition);

        List<Message> messages = topicMessages.get(partitionKey);
        if (messages == null || messages.isEmpty()) {
            return null;
        }

        long offset = consumerOffsets.getOrDefault(consumerKey, 0L);
        if (offset >= messages.size()) {
            return null;
        }

        Message message = messages.get((int) offset);
        consumerOffsets.put(consumerKey, offset + 1);
        return message;
    }

    /**
     * 检查主题是否存在（任何分区）
     */
    public boolean topicExists(String topic) {
        return topicMessages.keySet()
                .stream()
                .anyMatch(k -> k.startsWith(topic + "-"));
    }

    /**
     * 获取主题分区的写入偏移量（消息总数）
     */
    public long getWriteOffset(String topic, int partition) {
        List<Message> messages = topicMessages.get(getPartitionKey(topic, partition));
        return messages != null ? messages.size() : 0;
    }

    /**
     * 获取消费者的读取偏移量
     */
    public long getReadOffset(String consumerId, String topic, int partition) {
        return consumerOffsets.getOrDefault(
                getConsumerKey(consumerId, topic, partition),
                0L
        );
    }

    /**
     * 重置消费者偏移量
     */
    public void resetConsumerOffset(String consumerId, String topic, int partition, long offset) {
        consumerOffsets.put(
                getConsumerKey(consumerId, topic, partition),
                Math.max(0, offset)
        );
    }

    // 私有工具方法
    private String getPartitionKey(String topic, int partition) {
        return topic + "-" + partition;
    }

    private String getConsumerKey(String consumerId, String topic, int partition) {
        return consumerId + ":" + topic + ":" + partition;
    }

    /**
     * 清理过期消息（按配置的保留策略）
     */
    private void cleanExpiredMessages(String partitionKey) {
        List<Message> messages = topicMessages.get(partitionKey);
        if (messages == null || messages.size() <= config.getMaxMessagesPerTopic()) {
            return;
        }

        synchronized (messages) {
            // 清理超过数量限制的旧消息
            while (messages.size() > config.getMaxMessagesPerTopic()) {
                messages.remove(0);
            }
            // 注意：实际项目需要同步更新相关消费者偏移量
        }
    }
}
